#!/usr/bin/env python3

# Copyright (c) 2020-2021 Fpemud <fpemud@sina.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.


import os
import abc
import subprocess
from ._util import Util


class StorageLayout(abc.ABC):

    class MountResult(abc.ABC):

        def __init__(self, parent, mountDir, stage):
            self._parent = parent
            self._mountDir = mountDir
            self._stage = stage

        @property
        def stage(self):
            return self._stage

        def __enter__(self):
            return self

        def __exit__(self, type, value, traceback):
            self._parent.umount(self._mountDir)

    @staticmethod
    def get(storageLayoutName):
        if storageLayoutName == "bios-fat16":
            return StorageLayoutBiosFat16()
        elif storageLayoutName == "bios-ntfs":
            return StorageLayoutBiosNtfs()
        elif storageLayoutName == "efi-msr-ntfs":
            return StorageLayoutEfiMsrNtfs()
        else:
            assert False

    @abc.abstractmethod
    def initialize(self, diskFile, mountDir):
        pass

    @abc.abstractmethod
    def is_dirty(self, diskFile, mountDir):
        pass

    @abc.abstractmethod
    def enlarge_to(self, diskFile, mountDir, finalSize):
        pass

    @abc.abstractmethod
    def mount(self, diskFile, mountDir, stage):
        pass

    @abc.abstractmethod
    def is_mounted(self, mountDir, stage):
        pass

    @abc.abstractmethod
    def umount(self, mountDir):
        pass


class StorageLayoutBiosFat16(StorageLayout):

    _partitionLayout = (
        "msdos", [
            ("*", "fat16-lba-boot"),        # mkfs.vfat auto enable/disable LBA, we have no swith to control it
                                            # under this lba partition type, windows can do with both non-LBA and LBA filesystem
        ]
    )

    class MountResult(StorageLayout.MountResult):

        def __init__(self, parent, mountDir, stage, loop0, mnt0):
            super().__init__(parent, mountDir, stage)

            if stage == 1:
                assert loop0 is not None and mnt0 is None
            elif stage == 2:
                assert loop0 is not None and mnt0 is not None
            else:
                assert False
            self._loop0 = loop0
            self._mnt0 = mnt0

        @property
        def drive_c_loop_file(self):
            return self._loop0

        @property
        def drive_c_dir(self):
            assert self._mnt0 is not None
            return self._mnt0

    def initialize(self, diskFile, mountDir):
        dev, disk = Util.partedInitializeDiskFile(diskFile, self._partitionLayout)
        try:
            partiLoopFileList = Util.mountDiskFile(diskFile, self._partitionLayout, mountDir, partedDisk=disk)
            try:
                assert len(partiLoopFileList) == 1 and os.path.basename(partiLoopFileList[0]) == "partition0.loop"
                subprocess.check_call(["mkfs.vfat", "-F", "16", partiLoopFileList[0]], stdout=subprocess.DEVNULL)    # mkfs.vfat does not have a quiet option
            finally:
                Util.umountDiskFile(self._partitionLayout, mountDir)
        finally:
            # pyparted sucks: after guess-and-try, I think this is a usable way to CLEANLY do a clean-up
            del disk
            dev.getPedDevice().destroy()
            del dev

    def is_dirty(self, diskFile, mountDir):
        partiLoopFileList = Util.mountDiskFile(diskFile, self._partitionLayout, mountDir)
        try:
            assert len(partiLoopFileList) == 1 and os.path.basename(partiLoopFileList[0]) == "partition0.loop"
            return Util.fatCheckDirty(partiLoopFileList[0])
        finally:
            Util.umountDiskFile(self._partitionLayout, mountDir)

    def enlarge_to(self, diskFile, mountDir, finalSize):
        Util.partedEnlargeDiskFile(diskFile, self._partitionLayout, finalSize)

        partiLoopFileList = Util.mountDiskFile(diskFile, self._partitionLayout, mountDir)
        try:
            assert len(partiLoopFileList) == 1 and os.path.basename(partiLoopFileList[0]) == "partition0.loop"
            subprocess.check_call(["fatresize", "-s", "max", partiLoopFileList[0]])
        finally:
            Util.umountDiskFile(self._partitionLayout, mountDir)

    def mount(self, diskFile, mountDir, stage):
        partiLoopFileList = Util.mountDiskFile(diskFile, self._partitionLayout, mountDir)
        try:
            assert len(partiLoopFileList) == 1 and os.path.basename(partiLoopFileList[0]) == "partition0.loop"
            if stage == 1:
                return self.MountResult(self, mountDir, 1, partiLoopFileList[0], None)
            elif stage == 2 or stage == -1:
                ret = self.MountResult(self, mountDir, 2, partiLoopFileList[0], partiLoopFileList[0][:len(".loop") * -1] + ".mntdir")
                os.mkdir(ret.drive_c_dir)
                try:
                    subprocess.check_call(["fusefat", ret.drive_c_loop_file, ret.drive_c_dir, "-o", "rw+"],
                                          stdout=subprocess.DEVNULL,                       # fusefat has no quiet option
                                          stderr=subprocess.DEVNULL)                       # fusefat has no quiet option
                    return ret
                except BaseException:
                    os.rmdir(ret.drive_c_dir)
                    raise
            else:
                assert False
        except BaseException:
            Util.umountDiskFile(self._partitionLayout, mountDir)
            raise

    def is_mounted(self, mountDir, stage):
        assert stage in [1, 2, -1]

        partiLoopFile = os.path.join(mountDir, "partition0.loop")
        if os.path.exists(partiLoopFile) and Util.getMountDeviceForPath(partiLoopFile) == "fuseloop":
            if stage == 1:
                return True

        partiMountDir = os.path.join(mountDir, "partition0.mntdir")
        if os.path.exists(partiMountDir) and Util.getMountDeviceForPath(partiMountDir) == "fusefat":
            if stage == 2 or stage == -1:
                return True

        return False

    def umount(self, mountDir):
        partiMountDir = os.path.join(mountDir, "partition0.mntdir")
        if os.path.exists(partiMountDir):
            subprocess.run(["umount", partiMountDir])
            os.rmdir(partiMountDir)

        partiLoopFile = os.path.join(mountDir, "partition0.loop")
        if os.path.exists(partiLoopFile):
            Util.umountDiskFile(self._partitionLayout, mountDir)

    def check(self, diskFile):
        pass


class StorageLayoutBiosNtfs(StorageLayout):

    _partitionLayout = (
        "msdos", [
            ("*", "ntfs-boot"),
        ]
    )

    class MountResult(StorageLayout.MountResult):

        def __init__(self, parent, mountDir, stage, loop0, mnt0):
            super().__init__(parent, mountDir, stage)

            if stage == 1:
                assert loop0 is not None and mnt0 is None
            elif stage == 2:
                assert loop0 is not None and mnt0 is not None
            else:
                assert False
            self._loop0 = loop0
            self._mnt0 = mnt0

        @property
        def drive_c_loop_file(self):
            return self._loop0

        @property
        def drive_c_dir(self):
            assert self._mnt0 is not None
            return self._mnt0

    def initialize(self, diskFile, mountDir):
        dev, disk = Util.partedInitializeDiskFile(diskFile, self._partitionLayout)
        try:
            partiLoopFileList = Util.mountDiskFile(diskFile, self._partitionLayout, mountDir, partedDisk=disk)
            try:
                assert len(partiLoopFileList) == 1 and os.path.basename(partiLoopFileList[0]) == "partition0.loop"
                cmdList = [
                    "mkfs.ntfs",
                    "-s", "512",                            # specify sector size manually, to supress a warning
                    "-c", "4096",                           # specify cluster size manually, to supress a warning
                    "-H", "255", "-S", "63", "-p", "63",    # so that this partition is bootable by windows (63 sectors per track, partition starts from track 1)
                    "-F",
                    "-Q",
                    partiLoopFileList[0],
                ]
                subprocess.check_call(cmdList, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)    # quiet option of mkfs.ntfs is not totally quiet
            finally:
                Util.umountDiskFile(self._partitionLayout, mountDir)
        finally:
            # pyparted sucks: after guess-and-try, I think this is a usable way to CLEANLY do a clean-up
            del disk
            dev.getPedDevice().destroy()
            del dev

    def is_dirty(self, diskFile, mountDir):
        partiLoopFileList = Util.mountDiskFile(diskFile, self._partitionLayout, mountDir)
        try:
            assert len(partiLoopFileList) == 1 and os.path.basename(partiLoopFileList[0]) == "partition0.loop"
            return Util.ntfsCheckDirty(partiLoopFileList[0])
        finally:
            Util.umountDiskFile(self._partitionLayout, mountDir)

    def enlarge_to(self, diskFile, mountDir, finalSize):
        Util.partedEnlargeDiskFile(diskFile, self._partitionLayout, finalSize)

        partiLoopFileList = Util.mountDiskFile(diskFile, self._partitionLayout, mountDir)
        try:
            assert len(partiLoopFileList) == 1 and os.path.basename(partiLoopFileList[0]) == "partition0.loop"
            subprocess.run(["ntfsresize", partiLoopFileList[0]], input="y\n", text=True, check=True, stdout=subprocess.DEVNULL)
            subprocess.check_call(["ntfsfix", "-d", partiLoopFileList[0]], stdout=subprocess.DEVNULL)      # ntfsresize intentionally marks file system as dirty
        finally:
            Util.umountDiskFile(self._partitionLayout, mountDir)

    def mount(self, diskFile, mountDir, stage):
        partiLoopFileList = Util.mountDiskFile(diskFile, self._partitionLayout, mountDir, allowRoot=True)       # loop mount point must be accessible by root because ntfs-3g uses SUID to do fuse mount (why ntfs-3g uses SUID?)
        try:
            assert len(partiLoopFileList) == 1 and os.path.basename(partiLoopFileList[0]) == "partition0.loop"
            if stage == 1:
                return self.MountResult(self, mountDir, 1, partiLoopFileList[0], None)
            elif stage == 2 or stage == -1:
                ret = self.MountResult(self, mountDir, 2, partiLoopFileList[0], partiLoopFileList[0][:len(".loop") * -1] + ".mntdir")
                os.mkdir(ret.drive_c_dir)
                try:
                    subprocess.check_call(["ntfs-3g", ret.drive_c_loop_file, ret.drive_c_dir])
                    return ret
                except BaseException:
                    os.rmdir(ret.drive_c_dir)
                    raise
            else:
                assert False
        except BaseException:
            Util.umountDiskFile(self._partitionLayout, mountDir)
            raise

    def is_mounted(self, mountDir, stage):
        assert stage in [1, 2, -1]

        partiLoopFile = os.path.join(mountDir, "partition0.loop")
        if os.path.exists(partiLoopFile) and Util.getMountDeviceForPath(partiLoopFile) == "fuseloop":
            if stage == 1:
                return True

        partiMountDir = os.path.join(mountDir, "partition0.mntdir")
        if os.path.exists(partiMountDir) and Util.getMountDeviceForPath(partiMountDir) == "ntfs-3g":
            if stage == 2 or stage == -1:
                return True

        return False

    def umount(self, mountDir):
        partiMountDir = os.path.join(mountDir, "partition0.mntdir")
        if os.path.exists(partiMountDir):
            subprocess.run(["umount", partiMountDir])
            os.rmdir(partiMountDir)

        # FIXME: it seems ntfs-3g sucks that partition0.loop is still busy after umount command returns
        import time
        time.sleep(3)

        partiLoopFile = os.path.join(mountDir, "partition0.loop")
        if os.path.exists(partiLoopFile):
            Util.umountDiskFile(self._partitionLayout, mountDir)

    def check(self, diskFile):
        pass


class StorageLayoutEfiMsrNtfs(StorageLayout):

    # FIXME
    _partitionLayout = (
        "gpt", [
            ("100MiB", "esp"),
            ("128MiB", "msr"),
            ("*", "ntfs"),
        ]
    )

    class MountResult(StorageLayout.MountResult):

        def __init__(self, parent, mountDir, stage, loop0, loop1, loop2, mnt0, mnt2):
            super().__init__(parent, mountDir, stage)

            if stage == 1:
                assert loop0 is not None and mnt0 is None
                assert loop1 is not None
                assert loop2 is not None and mnt2 is None
            elif stage == 2:
                assert loop0 is not None and mnt0 is not None
                assert loop1 is not None
                assert loop2 is not None and mnt2 is not None
            else:
                assert False
            self._loop0 = loop0
            self._loop1 = loop1
            self._loop2 = loop2
            self._mnt0 = mnt0
            self._mnt2 = mnt2

        @property
        def esp_loop_file(self):
            return self._loop0

        @property
        def msr_loop_file(self):
            return self._loop1

        @property
        def drive_c_loop_file(self):
            return self._loop2

        @property
        def esp_dir(self):
            assert self._mnt0 is not None
            return self._mnt0

        @property
        def drive_c_dir(self):
            assert self._mnt2 is not None
            return self._mnt2

    def initialize(self, diskFile, mountDir):
        dev, disk = Util.partedInitializeDiskFile(diskFile, self._partitionLayout)
        try:
            partiLoopFileList = Util.mountDiskFile(diskFile, self._partitionLayout, mountDir, partedDisk=disk)
            try:
                assert len(partiLoopFileList) == 3

                # esp partition
                assert os.path.basename(partiLoopFileList[0]) == "partition0.loop"
                cmdList = [
                    "mkfs.vfat",
                    "-F", "32",
                    partiLoopFileList[0],
                ]
                subprocess.check_call(cmdList, stdout=subprocess.DEVNULL)    # mkfs.vfat does not have a quiet option

                # msr partition
                assert os.path.basename(partiLoopFileList[1]) == "partition1.loop"

                # sys partition
                assert os.path.basename(partiLoopFileList[2]) == "partition2.loop"
                cmdList = [
                    "mkfs.ntfs",
                    "-s", "512",                            # specify sector size manually, to supress a warning
                    "-c", "4096",                           # specify cluster size manually, to supress a warning
                    "-H", "255", "-S", "63", "-p", "63",    # so that this partition is bootable by windows (63 sectors per track, partition starts from track 1)
                    "-F",
                    "-Q",
                    partiLoopFileList[2],
                ]
                subprocess.check_call(cmdList, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)    # quiet option of mkfs.ntfs is not totally quiet
            finally:
                Util.umountDiskFile(self._partitionLayout, mountDir)
        finally:
            # pyparted sucks: after guess-and-try, I think this is a usable way to CLEANLY do a clean-up
            del disk
            dev.getPedDevice().destroy()
            del dev

    def is_dirty(self, diskFile, mountDir):
        partiLoopFileList = Util.mountDiskFile(diskFile, self._partitionLayout, mountDir)
        try:
            assert len(partiLoopFileList) == 3
            assert os.path.basename(partiLoopFileList[0]) == "partition0.loop"
            assert os.path.basename(partiLoopFileList[1]) == "partition1.loop"
            assert os.path.basename(partiLoopFileList[2]) == "partition2.loop"
            return Util.fatCheckDirty(partiLoopFileList[0]) or Util.ntfsCheckDirty(partiLoopFileList[2])
        finally:
            Util.umountDiskFile(self._partitionLayout, mountDir)

    def enlarge_to(self, diskFile, mountDir, finalSize):
        Util.partedEnlargeDiskFile(diskFile, self._partitionLayout, finalSize)

        partiLoopFileList = Util.mountDiskFile(diskFile, self._partitionLayout, mountDir)
        try:
            assert len(partiLoopFileList) == 3
            assert os.path.basename(partiLoopFileList[0]) == "partition0.loop"
            assert os.path.basename(partiLoopFileList[1]) == "partition1.loop"
            assert os.path.basename(partiLoopFileList[2]) == "partition2.loop"
            subprocess.run(["ntfsresize", partiLoopFileList[0]], input="y\n", text=True, check=True, stdout=subprocess.DEVNULL)
            subprocess.check_call(["ntfsfix", "-d", partiLoopFileList[0]], stdout=subprocess.DEVNULL)      # ntfsresize intentionally marks file system as dirty
        finally:
            Util.umountDiskFile(self._partitionLayout, mountDir)

    def mount(self, diskFile, mountDir, stage):
        partiLoopFileList = Util.mountDiskFile(diskFile, self._partitionLayout, mountDir, allowRoot=True)       # loop mount point must be accessible by root because ntfs-3g uses SUID to do fuse mount (why ntfs-3g uses SUID?)
        try:
            assert len(partiLoopFileList) == 1
            assert os.path.basename(partiLoopFileList[0]) == "partition0.loop"
            assert os.path.basename(partiLoopFileList[1]) == "partition1.loop"
            assert os.path.basename(partiLoopFileList[2]) == "partition2.loop"
            if stage == 1:
                return self.MountResult(self, mountDir, 1, partiLoopFileList[0], partiLoopFileList[1], partiLoopFileList[2], None, None)
            elif stage == 2 or stage == -1:
                ret = self.MountResult(self, mountDir, 1, partiLoopFileList[0], partiLoopFileList[1], partiLoopFileList[2],
                                       partiLoopFileList[0][:len(".loop") * -1] + ".mntdir",
                                       partiLoopFileList[2][:len(".loop") * -1] + ".mntdir")
                os.mkdir(ret.drive_c_dir)
                os.mkdir(ret.esp_dir)
                try:
                    subprocess.check_call(["ntfs-3g", ret.drive_c_loop_file, ret.drive_c_dir])
                    subprocess.check_call(["fusefat", ret.esp_loop_file, ret.esp_dir, "-o", "rw+"],
                                          stdout=subprocess.DEVNULL,                       # fusefat has no quiet option
                                          stderr=subprocess.DEVNULL)                       # fusefat has no quiet option
                    return ret
                except BaseException:
                    # FIXME: do force umount
                    Util.forceDelete(ret.esp_loop_file)
                    Util.forceDelete(ret.drive_c_loop_file)
                    raise
            else:
                assert False
        except BaseException:
            Util.umountDiskFile(self._partitionLayout, mountDir)
            raise

    def is_mounted(self, mountDir, stage):
        assert stage in [1, 2, -1]

        partiLoopFile = os.path.join(mountDir, "partition0.loop")
        if os.path.exists(partiLoopFile) and Util.getMountDeviceForPath(partiLoopFile) == "fuseloop":
            if stage == 1:
                return True

        partiMountDir0 = os.path.join(mountDir, "partition0.mntdir")
        partiMountDir2 = os.path.join(mountDir, "drive_c.mntdir")
        if os.path.exists(partiMountDir0) and Util.getMountDeviceForPath(partiMountDir0) == "ntfs-3g":
            if stage == 2 or stage == -1:
                return True
        if os.path.exists(partiMountDir2) and Util.getMountDeviceForPath(partiMountDir2) == "fusefat":
            if stage == 2 or stage == -1:
                return True

        return False

    def umount(self, mountDir):
        partiMountDir = os.path.join(mountDir, "partition0.mntdir")
        if os.path.exists(partiMountDir):
            subprocess.run(["umount", partiMountDir])
            os.rmdir(partiMountDir)
        partiMountDir = os.path.join(mountDir, "drive_c.mntdir")
        if os.path.exists(partiMountDir):
            subprocess.run(["umount", partiMountDir])
            os.rmdir(partiMountDir)

        partiLoopFile = os.path.join(mountDir, "partition0.loop")
        if os.path.exists(partiLoopFile):
            Util.umountDiskFile(self._partitionLayout, mountDir)
        partiLoopFile = os.path.join(mountDir, "partition2.loop")
        if os.path.exists(partiLoopFile):
            Util.umountDiskFile(self._partitionLayout, mountDir)

    def check(self, diskFile):
        pass
